# 导入LlamaIndex核心工具模块
from llama_index.core.tools import QueryEngineTool
# 导入摘要索引模块（用于生成文档摘要）
from llama_index.core import SummaryIndex
# 导入向量存储索引模块（用于语义搜索）
from llama_index.core import VectorStoreIndex
# 导入存储上下文模块（管理数据存储）
from llama_index.core import StorageContext
# 导入全局设置模块
from llama_index.core import Settings
# 导入简单目录读取器（用于加载文档）
from llama_index.core import SimpleDirectoryReader

# 加载文档：从"./data"目录读取所有文档
documents = SimpleDirectoryReader("./data").load_data()
# 初始化设置：设置文本块大小为1024 tokens
Settings.chunk_size = 1024
# 将文档解析为节点（Node）对象
nodes = Settings.node_parser.get_nodes_from_documents(documents)
# 初始化存储上下文（默认使用内存存储）
storage_context = StorageContext.from_defaults()
# 将解析后的节点添加到文档存储中
storage_context.docstore.add_documents(nodes)

# 创建摘要索引（用于生成文档摘要）
summary_index = SummaryIndex(nodes, storage_context=storage_context)
# 创建向量存储索引（用于语义搜索）
vector_index = VectorStoreIndex(nodes, storage_context=storage_context)

# 创建基于摘要索引的查询引擎，使用树形汇总响应模式和异步处理
list_query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize",  # 使用树形汇总方式生成回答
    use_async=True,  # 启用异步处理
)

# 创建基于向量索引的查询引擎（默认配置）
vector_query_engine = vector_index.as_query_engine()

# 创建摘要查询工具，设置描述信息
list_tool = QueryEngineTool.from_defaults(
    query_engine=list_query_engine,  # 绑定摘要查询引擎
    description=(  # 工具描述
        "Useful for summarization questions related to Paul Graham eassy on"
        " What I Worked On."  # 适用于Paul Graham《我做了什么》文章的摘要问题
    ),
)

# 创建向量查询工具，设置描述信息
vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,  # 绑定向量查询引擎
    description=(  # 工具描述
        "Useful for retrieving specific context from Paul Graham essay on What"
        " I Worked On."  # 适用于从Paul Graham文章中检索具体内容
    ),
)

# 导入路由查询引擎和LLM选择器
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# 创建路由查询引擎，自动选择最合适的查询工具
query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),  # 使用默认的LLM单选择器
    query_engine_tools=[  # 可选的查询工具列表
        list_tool,   # 摘要工具
        vector_tool, # 向量检索工具
    ],
)

# 执行第一个查询：获取文档摘要
response = query_engine.query("文档的汇总摘要是什么?")
# 打印查询结果
print(str(response))

# 执行第二个查询：询问Paul Graham在RICS之后的经历
response = query_engine.query("红楼梦的主要人物有哪些?")
# 打印查询结果
print(str(response))